# 识别多边形
import cv2
import math
import numpy as np

# 创建窗口
windowname = 'threat'
cv2.namedWindow(windowname)


def p(x):
    pass


# 创建滑动条
cv2.createTrackbar('threshold', windowname, 250, 255, p)

COEFFICIENT = 0.02


class ShapeDetector:
    # 初始化类
    def __init__(self):
        # 字典类型对应每一种图形的计数器
        self.counter = {"unrecognized": 0, "triangle": 0, "rhombus": 0, "rectangle": 0, "pentagon": 0,
                        "hexagon": 0, "circle": 0}
        # 初始化图形类型为不可识别
        self.shape = "unrecognized image"
        # 图形顶点集置空
        self.approx = []
        # 初始化该图形的周长为0
        self.peri = 0

    # 计算欧式距离(主要作用通过计算边长区分菱形和长方形)
    def distance(self, x1, y1, x2, y2):
        return math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)

    def detect(self, c, frame):
        # cv2.arcLength函数返回周长
        self.peri = cv2.arcLength(c, True)

        # cv2.approxPolyDP用多边形取拟合，返回的是顶点的列表
        self.approx = cv2.approxPolyDP(c, COEFFICIENT * self.peri, True)
        # self.approx = cv2.approxPolyDP(c, 0.2, True)
        d = len(self.approx)
        for i in range(d):
            frame2 = cv2.line(frame, (self.approx[i][0][0], self.approx[i][0][1]),
                             (self.approx[(i + 1) % d][0][0], self.approx[(i + 1) % d][0][1]), (255, 0, 0), 3)
            cv2.circle(frame2, (self.approx[i][0][0], self.approx[i][0][1]), 4, (0, 0, 255), -1)

        # 3个顶点，三角形
        if d == 3:
            self.shape = "triangle"
        # 同理，四个顶点，四边形
        elif d == 4:
            # 计算相邻两边的长度，做差判在误差范围内是否相等
            dist1 = self.distance(self.approx[0][0][0], self.approx[0][0][1], self.approx[1][0][0],
                                  self.approx[1][0][1])
            dist2 = self.distance(self.approx[0][0][0], self.approx[0][0][1], self.approx[3][0][0],
                                  self.approx[3][0][1])
            result = math.fabs(dist1 - dist2)
            # print(result)
            # 误差小于10，可近似认为相等，为菱形
            if result <= 10:
                self.shape = "rhombus"
            else:
                self.shape = "rectangle"
        # 五边形
        elif d == 5:
            self.shape = "pentagon"
        # 六边形
        elif d == 6:
            self.shape = "hexagon"
        # 圆
        else:
            self.shape = "unrecognized"

        # # 相应形状计数器加一
        # self.counter[self.shape] += 1

        cv2.imshow(windowname, frame2)

        # 返回形状
        return self.shape, len(self.approx)

    # def Display(self):
    #     # 展现结果
    #     for kind in self.counter.keys():
    #         # if self.counter[kind] != 0:
    #         print("The number of {} is {}".format(kind, self.counter[kind]))


def main():
    o_shape = 'unrecognized'
    cap = cv2.VideoCapture(0)
    while cap.isOpened():
        ret, frame = cap.read()
        cv2.imshow('frame', frame)

        # 将图片转换为灰度图片
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        # 高斯滤波,图像平滑处理
        blurred = cv2.GaussianBlur(gray, (3, 3), 0)
        # 根据阈值，将灰度图片转化为黑白两色图片
        threshold = cv2.getTrackbarPos('threshold', windowname)
        thresh = cv2.threshold(blurred, threshold, 255, cv2.THRESH_BINARY)[1]

        # 返回图片和图中轮廓信息，查找最大轮廓返回到cnts中
        cnts = cv2.findContours(thresh.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[-2]
        if len(cnts):
            area = []
            for i in range(len(cnts)):
                area.append(cv2.contourArea(cnts[i]))
                max_idx = np.argmax(np.array(area))


            # 创建一个识别器实例
            sd = ShapeDetector()
            # 对轮廓进行处理
            n_shape, point = sd.detect(cnts[max_idx], frame)

            if n_shape != o_shape:
                print(n_shape)
                print(point)
                o_shape = n_shape

        # # 输出结果
        # sd.Display()
        #
        # cv2.imshow(windowname, thresh)

        c = cv2.waitKey(1)
        if c == ord('q'):
            break
    cap.release()
    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()
